package com.campus.websocket;

import com.campus.service.AIService;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * AI对话，流式返回数据的 webSocket
 *
 * @author TaiZai
 * @data 2024/11/10 16:51
 */
@Component
@ServerEndpoint("/ws/ai-chat/{uid}")
public class AIChatWebSocketServer {

    private static final Map<String, Session> sessionMap = new HashMap<>(); // 存放会话对象

    private static AIService aiService;

    @Autowired
    public void setAiService(AIService aiService) {
        AIChatWebSocketServer.aiService = aiService;
    }
    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("uid") String uid) {
        System.out.println("客户端：" + uid + " 已连接");
        session.setMaxTextMessageBufferSize(1024 * 1024);
        session.setMaxBinaryMessageBufferSize(1024 * 1024);
        sessionMap.put(uid, session);
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message, @PathParam("uid") String uid) {
         System.out.println("收到来自客户端：" + uid + "的信息:" + message);
        if (aiService != null) {
            Session session = sessionMap.get(uid);
            if (session != null && session.isOpen()) {
                aiService.chatWithCloseAI(message, session); // 调用服务以流式传输消息
            }
        } else {
            System.out.println("AIService 未注入！");
        }
    }

    /**
     * 连接关闭调用的方法
     *
     * @param uid
     */
    @OnClose
    public void onClose(@PathParam("uid") String uid) {
        System.out.println("连接断开:" + uid);
        sessionMap.remove(uid);
    }

    /**
     * 向特定客户端发送消息
     */
    public void sendToClient(String uid, String message) {
        Session session = sessionMap.get(uid);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 群发
     *
     * @param message
     */
    public void sendToAllClient(String message) {
        Collection<Session> sessions = sessionMap.values();
        for (Session session : sessions) {
            try {
                // 服务器向客户端发送消息
                session.getBasicRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
